
import unittest
from tempfile import NamedTemporaryFile

from forpp import TextBlock, Parser, FunctionBlock, IfBlock, ForppException, ForBlock, MacroBlock

class TestReplace(unittest.TestCase):

    def test_expand_simple_variable(self):
        string = [
            '<FOO>'
        ]
        exp = [
            'hello'
        ]
        block = TextBlock(string, macros={'FOO': MacroBlock('FOO', lines=['hello'])})
        res = block.expand()
        self.assertEqual(res, exp)
        
    def test_expand_expression(self):
        string = [
            '<FOO+1>',
            '<FOO-2>',
            '<FOO/3>',
            '<FOO*4>',
        ]
        exp = [
            '2',
            '-1',
            '0',
            '4',
        ]
        block = TextBlock(string, macros={'FOO': MacroBlock('FOO', lines=['1'])})
        res = block.expand()
        self.assertEqual(res, exp)
        
    def test_expand_function_with_one_arg(self):
        string = [
            'FOO( 1 )'
        ]
        exp = [
            '2'
        ]
        block = TextBlock(string, macros={'FOO': FunctionBlock('FOO', ['ARG1'], lines=['<ARG1+1>'])})
        res = block.expand()
        self.assertEqual(res, exp)
        
    def test_expand_function_with_multiple_args(self):
        string = [
            'FOO( integer, 2, var )'
        ]
        exp = [
            'integer, dimension(2) :: var'
        ]
        args = ['TYPE', 'DIM', 'NAME']
        fcn_lines = ['<TYPE>, dimension(<DIM>) :: <NAME>']
        block = TextBlock(string, macros={'FOO': FunctionBlock('FOO', args, lines=fcn_lines)})
        res = block.expand()
        self.assertEqual(res, exp)
        
    def test_expand_function_with_args_with_braces(self):
        string = [
            'FOO( integer(kind=4), 2, var )'
        ]
        exp = [
            'integer(kind=4), dimension(2) :: var'
        ]
        args = ['TYPE', 'DIM', 'NAME']
        fcn_lines = ['<TYPE>, dimension(<DIM>) :: <NAME>']
        block = TextBlock(string, macros={'FOO': FunctionBlock('FOO', args, lines=fcn_lines)})
        res = block.expand()
        self.assertEqual(res, exp)

    def test_expand_if_clause(self):
        tests = ['<DIM> == 0', '<DIM> == 1', None]
        test_blocks = [
            ['if branch'],
            ['elseif branch'],
            ['else branch']
        ]
        if_block = IfBlock(tests, test_blocks, macros={'DIM': MacroBlock('DIM', lines=['0'])})
        string = [
            if_block
        ]
        exp = [
            'if branch'
        ]
        block = TextBlock(string)
        res = block.expand()
        self.assertEqual(res, exp)
        
    def test_expand_elseif_clause(self):
        tests = ['<DIM> == 0', '<DIM> == 1', None]
        test_blocks = [
            ['if branch'],
            ['elseif branch'],
            ['else branch']
        ]
        if_block = IfBlock(tests, test_blocks, macros={'DIM': TextBlock('1')})
        string = [
            if_block
        ]
        exp = [
            'elseif branch'
        ]
        block = TextBlock(string)
        res = block.expand()
        self.assertEqual(res, exp)
        
    def test_expand_else_clause(self):
        tests = ['<DIM> == 0', '<DIM> == 1', None]
        test_blocks = [
            ['if branch'],
            ['elseif branch'],
            ['else branch']
        ]
        if_block = IfBlock(tests, test_blocks, macros={'DIM': TextBlock('3')})
        string = [
            if_block
        ]
        exp = [
            'else branch'
        ]
        block = TextBlock(string)
        res = block.expand()
        self.assertEqual(res, exp)
        
    def test_expand_for_block(self):
        lines = [
            'integer :: var<VAR>'
        ]
        exp = [
            'integer :: var1',
            'integer :: var2',
            'integer :: var3',
        ]
        block = ForBlock('VAR', '[1, 2, 3]', lines=lines)
        res = block.expand()
        self.assertEqual(res, exp)

    def test_expand_for_block_with_nested_function(self):
        lines = [
            'FOO( <VAR> )'
        ]
        exp = [
            'integer :: var1',
            'integer :: var2',
            'integer :: var3',
        ]
        macros = {'FOO': FunctionBlock('FOO', ['ARG1'], lines=['integer :: var<ARG1>'])}
        block = ForBlock('VAR', '[1, 2, 3]', lines=lines, macros=macros)
        res = block.expand()
        self.assertEqual(res, exp)

class TestParse(unittest.TestCase):

    def test_create_new_parser(self):
        parser = Parser(None)

    def test_parse_empty_string(self):
        parser = Parser(None)
        res = parser.parse(["\n"])
        exp = [
            '\n'
        ]
        self.assertEqual(parser.macros, {})
        self.assertEqual(res, exp)

    def test_parse_simple_single_line_define(self):
        parser = Parser(None)
        res = parser.parse(["!#DEFINE NAME Value\n"])
        exp = []
        self.assertEqual(parser.macros.keys(), ['NAME'])
        self.assertEqual(parser.macros['NAME'].Type, 'DEFINE')
        self.assertEqual(res, exp)
    
    def test_parse_define_block(self):
        parser = Parser(None)
        res = parser.parse([
            '!#DEFINE NAME\n',
            'Value\n',
            '!#END DEFINE',
        ])
        exp = []
        self.assertEqual(parser.macros.keys(), ['NAME'])
        self.assertEqual(parser.macros['NAME'].Type, 'DEFINE')
        self.assertEqual(res, exp)
    
    def test_parse_non_empty_strign(self):
        parser = Parser(None)
        res = parser.parse(["module main\n"])
        exp = [
            'module main\n'
        ]
        self.assertEqual(parser.macros, {})
        self.assertEqual(res, exp)
    
    def test_multi_line_simple_define(self):
        parser = Parser(None)
        res = parser.parse([
            "!#DEFINE NAME Some \\\n",
            "!             multi line \\\n",
            "!             define\n"
        ])
        exp = []
        self.assertEqual(parser.macros.keys(), ['NAME'])
        self.assertEqual(parser.macros['NAME'].Type, 'DEFINE')
        self.assertEqual(res, exp)
    
    def test_for_block(self):
        parser = Parser(None)
        res = parser.parse([
            "!#FOR VAR IN [1, 2, 3]\n",
            "   integer :: var<VAR>\n",
            "!#END FOR\n"
        ])
        self.assertEqual(res[0].Type, 'FOR')
        self.assertEqual(res[0].var, 'VAR')
        self.assertEqual(res[0]._list, ['1', '2', '3'])
        
    def test_parse_define_containing_block(self):
        parser = Parser(None)
        res = parser.parse([
            "!#DEFINE BEGIN_SUITE( NAME, USES )",
            "module test_<NAME>_mod",
            "  use pFUnit",
            "!# FOR USE IN <USES>",
            "  use <USE>",
            "!# END FOR",
            "  implicit none",
            "!#END DEFINE"
        ])
        exp = []
        self.assertEqual(res, [])
        self.assertEqual(parser.macros.keys(), ['BEGIN_SUITE'])
        self.assertEqual(parser.macros['BEGIN_SUITE']._lines, [
            "module test_<NAME>_mod",
            "  use pFUnit",
            "!# FOR USE IN <USES>",
            "  use <USE>",
            "!# END FOR",
            "  implicit none",
        ])
    
    def test_nested_for_block(self):
        parser = Parser(None)
        res = parser.parse([
            "!#FOR VAR IN [1,2,3]",
            "!# IF <VAR> == 0",
            "  integer :: var",
            "!# ELSE",
            "  integer, dimension(<VAR>) :: var",
            "!# END IF",
            "  integer :: i",
            "!#END FOR"
        ])
        exp = [
            "!# IF <VAR> == 0",
            "  integer :: var",
            "!# ELSE",
            "  integer, dimension(<VAR>) :: var",
            "!# END IF",
            "  integer :: i",
        ]
        self.assertEqual(res[0].Type, 'FOR')
        self.assertEqual(res[0].var, 'VAR')
        self.assertEqual(res[0]._list, ['1','2','3'])
        self.assertEqual(res[0]._lines, exp)
        
    def test_parse_simple_if_block(self):
        parser = Parser(None)
        res = parser.parse([
            "!#IF <VAR> == 0",
            "  if block",
            "!#END IF",
        ])
        exp = [
            "  if block",
        ]
        self.assertEqual(res[0].Type, 'IF')
        self.assertEqual(res[0]._tests, ['<VAR> == 0'])
        self.assertEqual(res[0]._test_lines, {
            '<VAR> == 0': ['  if block'],
        })
        
    def test_parse_if_elseif_block(self):
        parser = Parser(None)
        res = parser.parse([
            "!#IF <VAR> == 0",
            "  if block",
            "!#ELSEIF <VAR> == 1",
            "  elseif block",
            "!#END IF",
        ])
        exp = [
            "  if block",
        ]
        self.assertEqual(res[0].Type, 'IF')
        self.assertEqual(res[0]._tests, ['<VAR> == 0', '<VAR> == 1'])
        self.assertEqual(res[0]._test_lines, {
            '<VAR> == 0': ['  if block'],
            '<VAR> == 1': ['  elseif block'],
        })
        
    def test_parse_if_elseif_else_block(self):
        parser = Parser(None)
        res = parser.parse([
            "!#IF <VAR> == 0",
            "  if block",
            "!#ELSEIF <VAR> == 1",
            "  elseif block",
            "!#ELSE",
            "  else block",
            "!#END IF",
        ])
        exp = [
            "  if block",
        ]
        self.assertEqual(res[0].Type, 'IF')
        self.assertEqual(res[0]._tests, ['<VAR> == 0', '<VAR> == 1', None])
        self.assertEqual(res[0]._test_lines, {
            '<VAR> == 0': ['  if block'],
            '<VAR> == 1': ['  elseif block'],
            None: ['  else block'],
        })
        
    def test_parse_include(self):
        parser = Parser(None)
        tmpfile = NamedTemporaryFile(delete=False)
        tmpfile.writelines([
            "!#DEFINE MACRO \\\n",
            "!    macro value\n",
            "!#IF <VAR> == 0\n",
            "  if block\n",
            "!#END IF\n",
        ])
        tmpfile.close()
        try:
            res = parser.parse([
                "!#INCLUDE %s" % tmpfile.name,
            ])
            exp = [
                "  if block",
            ]
            self.assertEqual(res[0].Type, 'IF')
            self.assertEqual(parser.macros.keys(), ['MACRO'])
        finally:
            import os
            os.remove(os.path.abspath(tmpfile.name))
        
    def test_parse_set(self):
        parent = TextBlock()
        block = TextBlock(parent=parent)
        parser = Parser(block)
        res = parser.parse([
            "!#SET VAR Value",
        ])
        self.assertEqual(res, [])
        self.assertEqual(parser.macros, {})
        self.assertEqual(block._macros, {})
        self.assertEqual(parent._macros.keys(), ['VAR'])
        self.assertEqual(parent._macros['VAR']._lines, ['Value'])
        
    def test_parse_append_to_empty_list(self):
        parent = TextBlock(macros={'VAR': '[]'})
        block = TextBlock(parent=parent)
        parser = Parser(block)
        res = parser.parse([
            "!#APPEND VAR Value",
        ])
        self.assertEqual(res, [])
        self.assertEqual(parser.macros, {})
        self.assertEqual(block._macros, {})
        self.assertEqual(parent._macros.keys(), ['VAR'])
        self.assertEqual(parent._macros['VAR']._lines, ['[Value]'])
        
    def test_parse_append_to_list_of_one_elemnent(self):
        parent = TextBlock(macros={'VAR': '[Value]'})
        block = TextBlock(parent=parent)
        parser = Parser(block)
        res = parser.parse([
            "!#APPEND VAR AnotherValue",
        ])
        self.assertEqual(res, [])
        self.assertEqual(parser.macros, {})
        self.assertEqual(block._macros, {})
        self.assertEqual(parent._macros.keys(), ['VAR'])
        self.assertEqual(parent._macros['VAR']._lines, ['[Value, AnotherValue]'])
    
    def test_parse_not_terminated_define_block(self):
        parser = Parser(None)
        self.assertRaises(ForppException,
            lambda: parser.parse([
                '!#DEFINE NAME\n',
                'Value\n',
            ]))

if __name__ == '__main__':
    unittest.main()
